#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <chrono>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>

#define BUFFER_SIZE 1024
#define DEFAULT_PORT 8888

using namespace std;
using namespace std::chrono;

void handle_tcp_connection(int client_socket)
{
    char buffer[BUFFER_SIZE];
    string filename = "received_file_tcp";
    
    ofstream outfile(filename, ios::binary);
    if (!outfile)
	{
        cerr << "Failed to create file: " << filename << endl;
        return;
    }

    auto start = high_resolution_clock::now();
    ssize_t total_bytes = 0;

    while (true)
	{
        ssize_t bytes_received = recv(client_socket, buffer, BUFFER_SIZE, 0);
        if (bytes_received <= 0) break;
        
        outfile.write(buffer, bytes_received);
        total_bytes += bytes_received;
    }

    auto end = high_resolution_clock::now();
    auto duration = duration_cast<milliseconds>(end - start);
    double rate = (total_bytes / (1024.0 * 1024.0)) / (duration.count() / 1000.0);

    cout << "TCP Transfer completed. " << total_bytes << " bytes received." << endl;
    cout << "Time: " << duration.count() << " ms" << endl;
    cout << "Transfer rate: " << rate << " MB/s" << endl;

    outfile.close();
    close(client_socket);
}

void handle_udp_connection(int server_socket)
{
    char buffer[BUFFER_SIZE];
    sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    
    string filename = "received_file_udp";
    ofstream outfile(filename, ios::binary);
    if (!outfile)
	{
        cerr << "Failed to create file: " << filename << endl;
        return;
    }

    auto start = high_resolution_clock::now();
    ssize_t total_bytes = 0;

    while (true)
	{
        ssize_t bytes_received = recvfrom(server_socket, buffer, BUFFER_SIZE, 0, (struct sockaddr*)&client_addr, &client_len);
        if (bytes_received <= 0) break;
        
        //确认信息
        if (bytes_received == 5 && strncmp(buffer, "END", 3) == 0) break;
        
        outfile.write(buffer, bytes_received);
        total_bytes += bytes_received;
    }

    auto end = high_resolution_clock::now();
    auto duration = duration_cast<milliseconds>(end - start);
    double rate = (total_bytes / (1024.0 * 1024.0)) / (duration.count() / 1000.0);

    cout << "UDP Transfer completed. " << total_bytes << " bytes received." << endl;
    cout << "Time: " << duration.count() << " ms" << endl;
    cout << "Transfer rate: " << rate << " MB/s" << endl;

    outfile.close();
}

int main(int argc, char* argv[])
{
    if (argc < 2) {
        cerr << "Usage: " << argv[0] << " <tcp|udp>" << endl;
        return 1;
    }

    string protocol = argv[1];
    int server_socket;
    
    if (protocol == "tcp")
	{
        //TCP启用
        server_socket = socket(AF_INET, SOCK_STREAM, 0);
        if (server_socket < 0)
		{
            perror("TCP socket creation failed");
            return 1;
        }

        sockaddr_in server_addr;
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(DEFAULT_PORT);

        if (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr))
		{
            perror("TCP bind failed");
            return 1;
        }

        listen(server_socket, 5);
        cout << "TCP Server listening on port " << DEFAULT_PORT << endl;

        sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);
        int client_socket = accept(server_socket, (struct sockaddr*)&client_addr, &client_len);
        if (client_socket < 0)
		{
            perror("TCP accept failed");
            return 1;
        }

        cout << "TCP Connection accepted from " << inet_ntoa(client_addr.sin_addr) << endl;
        handle_tcp_connection(client_socket);
        
    }
	else if (protocol == "udp")
	{
        //UDP启用
        server_socket = socket(AF_INET, SOCK_DGRAM, 0);
        if (server_socket < 0)
		{
            perror("UDP socket creation failed");
            return 1;
        }

        sockaddr_in server_addr;
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(DEFAULT_PORT);

        if (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)))
		{
            perror("UDP bind failed");
            return 1;
        }

        cout << "UDP Server listening on port " << DEFAULT_PORT << endl;
        handle_udp_connection(server_socket);
        
    }
	else
	{
        cerr << "Invalid protocol. Use 'tcp' or 'udp'." << endl;
        return 1;
    }

    close(server_socket);
    return 0;
}